## Calibrate eta/rho for each term structure

# Create functions to return term structure for a given rho/eta pair
findrho = function(eta, g=global_growth_percap, r.target, n.yrs=10) {
  # pick rho such that on average the CE rate in the first n.yrs
  # equals r.target (e.g., 3%).
  # r.target = -1/t * log(E[exp(-sum_t(r))])
  #      = -1/t * log(E[exp(-sum_t(rho+eta*g))])
  #      = -1/t * log(exp(-t*rho*E[exp(-eta*sum_t(g))])
  #      = -1/t * [ log(exp(-t*rho)) + log(E[exp(-eta*sum_t(g))]) ]
  #      = -1/t * [ -t*rho + log(E[exp(-eta*sum_t(g))]) ]
  #      = rho - 1/t*log(E[exp(-eta*sum_t(g))]) 
  # or: rho = r.target -  -1/t*log(E[exp(-eta*sum_t(g))]) 
  disc.facts.g = exp(-eta*apply(g[1:n.yrs,], 2, cumsum)) # n.yrs x B
  e.df.g = apply(disc.facts.g, 1, mean) # expected discount factor for each year
  r.ce.g = log(e.df.g) * -1/c(1:n.yrs)
  rho = r.target - r.ce.g # for each year, find the rho that works
  rho = mean(rho)
  return(rho)
}


# Function to calculate the fitted CE term structure for any eta/rho pair
calc_scc_ce = function(eta, rho=NULL, g=global_growth_percap, return.items='all', n.yrs=10,
                       B=ncol(g), target.ce) {
  # This function calculates the CE term structure for a given eta and target rate,
  # It also calculates rho (consistent with near-term target rate), 
  if (is.null(rho)) rho = findrho(eta, g=g, n.yrs=n.yrs, r.target=round(target.ce[1], 3))

  # Compute r = rho + eta*g for each period and state
  r.temp = rho + eta*g
  
  # Compute discount factor for each year and draw (cumsum discounts for each year through t)
  disc.factors = exp(-apply(r.temp, 2, cumsum)) # for each year and draw
  r.ce.path = -log(apply(disc.factors, 1, mean))/1:nrow(disc.factors)
  names(r.ce.path) = rownames(disc.factors)

  if (return.items=='rce') return(r.ce=r.ce.path)
  if (return.items=='all') {
    return(list(r.ce=data.table(year=rownames(g), r.ce=r.ce.path, r.ce.target=target.ce[1:nrow(g)]),
                rho=rho, eta=eta))
  }
}

calc_scc_ce(eta=0, target.ce=ce_rates[,'RW','3%']) # test. rho=3%, eta=0
calc_scc_ce(eta=0.5, target.ce=ce_rates[,'RW','3%']) # test. rho=1.7%, eta=0.5
calc_scc_ce(eta=1, target.ce=ce_rates[,'RW','3%']) # test. rho=0.45%, eta=1


# Function to calculate the mean square distance between the target and fitted rate, for a given eta
NP_MWS_ce_dist = function(eta, target.ce, n.yrs=10, match.rate=TRUE, ...) {
  # This function calculates the squared distance (deviation) between the target term CE structure
  # and the term structure calculated using eta, applied to the growth distribution (global_growth_percap)
  # eta...scalar containing eta (rho will be derived)
  # w...tt vector containing weights for computing sum of squares
  
  ce.temp = calc_scc_ce(eta, return.items='rce', n.yrs=n.yrs, target.ce=target.ce, ...)
  
  if (match.rate==TRUE) {
    mse = mean( ( ce.temp - target.ce[1:length(ce.temp)] )^2 )
  } else {
    DF.fit = exp(-c(1:length(ce.temp))*ce.temp)
    DF.target = exp(-c(1:length(ce.temp))*(target.ce[1:length(ce.temp)]))
    mse = mean( ( DF.fit - DF.target )^2 )
  }
  return(mse)
}

output <- array(list(), dim=c(dim(ce_rates)[2:3]), 
                dimnames=dimnames(ce_rates)[2:3])

set.seed(1)

library(progress)
pb = progress_bar$new(total=prod(dim(output)))

n.yrs = 10
st = Sys.time()
for (m in 1:dim(output)[1]) {
  for (r in 1:dim(output)[2]) {
    
    # Solve for optimal eta
    sol.temp = nlm(NP_MWS_ce_dist, target.ce=ce_rates[,m,r], 1.3, 
                   n.yrs=n.yrs, iterlim=1000, gradtol=1e-6,
                   g=global_growth_percap)
    
    if (sol.temp$code>=4) {message(paste(m,r,' failed')); next}
    
    # Output term structure for optimal eta
    ce.fit = calc_scc_ce(sol.temp$estimate, target.ce=ce_rates[,m,r], n.yrs=n.yrs,
                         g=global_growth_percap)
    
    # Find eta when rho is truncated at zero
    if (ce.fit$rho<0) {
      unconstrained.rho = ce.fit$rho
      unconstrained.eta = ce.fit$eta
      
      eta.grid.temp = seq(0,2, by=0.01)
      
      near.term.diff = function(eta) {
        near.term.g.ce = log(apply(exp(-eta*apply(global_growth_percap[1:n.yrs,], 2, cumsum)), 1, mean)) * -1/c(1:n.yrs)
        return(mean((near.term.g.ce - ce_rates[1:n.yrs,m,r])^2))
      }
      ntd = sapply(eta.grid.temp, near.term.diff)
      cbind(eta.grid.temp, ntd)
      plot(ntd ~ eta.grid.temp); abline(h=0)
      eta.opt = eta.grid.temp[which.min(ntd)]
      
      ce.fit = calc_scc_ce(eta.opt,rho=0, target.ce=ce_rates[,m,r], n.yrs=n.yrs,
                           g=global_growth_percap)
      ce.fit$unconstrained.rho = unconstrained.rho
      ce.fit$unconstrained.eta = unconstrained.eta
      
    }
    
    ce.fit$sse = sol.temp$`minimum` # save sum of squared errors
    
    output[[m,r]] = ce.fit
    pb$tick()
  }
  message(paste('\n',m,'of',dim(output)[1],' models done'))
}
ed = Sys.time()
difftime(ed, st)

res = sapply(output['BR',], function(x) c(x$rho, x$eta))

output.3pct = output[['BR','3%']]

rho = output.3pct$rho
eta = output.3pct$eta


### Plot fitted vs target rates
# A function to plot fitted vs target term structure
plot.ces = function(x, addline=FALSE, main='', maxbuff = 0.01, lty.fit=2, ...) {
  if (addline==FALSE) {
    plot(r.ce.target ~ year, data=x$`r.ce`, lty=1, lwd=3, type='l', ylim=c(0, max(r.ce.target)+maxbuff),
         xlim=range(as.numeric(year)), main=main, cex.main=1.5, ylab='Discount Rate from Year to 2020', xlab='Year', cex.lab=1.5, cex.axis=1.5, ...)
    axis(1, cex.axis=1.5, at=x$`r.ce`[, min(year)]); abline(v=x$`r.ce`[, min(year)], lty=3, col='gray70')
    grid(col='gray70')
  } else{
    lines(r.ce.target ~ year, data=x$`r.ce`, lty=1, lwd=3,  ...)
  }
  lines(r.ce ~ year, data=x$`r.ce`, lty=lty.fit, lwd=3, ...)
}

## BR term structure
pdf(figure.dir%&%'/BauerRudebusch_RhoEta_Estimates_'%&%today()%&%'.pdf',
    family='serif', width=7*1.5, height=7)
dimnames(output)

plot.ces(output['BR','5%'][[1]], col=3, main='Bauer Rudebusch (2020) Fits', maxbuff=0.01)
plot.ces(output['BR','3%'][[1]], addline=TRUE, col=1, cex.axis=1.5, cex.lab=2)
plot.ces(output['BR','2%'][[1]], addline=TRUE, col=4, cex.axis=1.5, cex.lab=2)
plot.ces(output['BR','1.5%'][[1]], addline=TRUE, col=2, cex.axis=1.5, cex.lab=2)
abline(h=0)

legend.values = sapply(output['BR', c('5%','3%','2%','1.5%')], function(x) c(rho=x$rho, eta=x$eta))
constr.rho = which(legend.values['rho',]==0)
legend.values['rho',] = round(legend.values['rho',],3)*100
legend.values['eta',] = round(legend.values['eta',],2)
legend.values['rho',] = paste0(legend.values['rho',],'%')
legend.values[,constr.rho] = paste0(legend.values[,constr.rho],'*')

legend('topright', lty=2:1, lwd=3, bg='white', legend=c('Fitted Rate', 'Target Rate'), cex=1.3, col='black')
legend('top', fill=c(3,1,4,2), border=NA, bg='white', cex=1.15,
       legend = c(
         as.expression(bquote(
           paste(r[0], ' = 5%  ',
                 rho,' = ', .(legend.values['rho','5%']), '    ',
                 eta,' = ', .(legend.values['eta','5%'])))),
         as.expression(bquote(
           paste(r[0], ' = 3%  ',
                 rho,' = ', .(legend.values['rho','3%']), '    ',
                 eta,' = ', .(legend.values['eta','3%'])))),
         as.expression(bquote(
           paste(r[0], ' = 2%  ',
                 rho,' = ', .(legend.values['rho','2%']), '    ',
                 eta,' = ', .(legend.values['eta','2%'])))),
         as.expression(bquote(
           paste(r[0], ' = 1.5%  ',
                 rho,' = ', .(legend.values['rho','1.5%']), '    ',
                 eta,' = ', .(legend.values['eta','1.5%']))))
       )
)
dev.off()


### ALL target term structures
{
  # dev.off()
  pdf(figure.dir%&%'/All_RhoEta_Estimates_'%&%today()%&%'.pdf',
      family='serif', width=10,height=12) #width=10, height=7)
  par(mfrow=c(3,2))
  # B&R
  plot.ces(output['BR','5%'][[1]], maxbuff=0.04, col=3, main='Bauer Rudebusch (2020) Fits')
  plot.ces(output['BR','3%'][[1]], addline=TRUE, col=1)
  plot.ces(output['BR','2%'][[1]], addline=TRUE, col=4)
  plot.ces(output['BR','1.5%'][[1]], addline=TRUE, col=2)
  
  abline(h=0)
  legend.values = sapply(output['BR', c('5%','3%','2%','1.5%')], function(x) c(rho=x$rho, eta=x$eta))
  constr.rho = which(legend.values['rho',]==0)
  legend.values['rho',] = round(legend.values['rho',],3)*100
  legend.values['eta',] = round(legend.values['eta',],2)
  legend.values['rho',] = paste0(legend.values['rho',],'%')
  legend.values[,constr.rho] = paste0(legend.values[,constr.rho],'*')
  
  legend('right', lty=2:1, lwd=3, bg='white', legend=c('Fitted Rate', 'Target Rate'), cex=1.4, col='black')
  legend('top', fill=c(3,1,4,2), border=NA, bg='white', cex=1.5,
         legend = c(
           as.expression(bquote(
             paste(r[0], ' = 5%  ',
                   rho,' = ', .(legend.values['rho','5%']), '    ',
                   eta,' = ', .(legend.values['eta','5%'])))),
           as.expression(bquote(
             paste(r[0], ' = 3%  ',
                   rho,' = ', .(legend.values['rho','3%']), '    ',
                   eta,' = ', .(legend.values['eta','3%'])))),
           as.expression(bquote(
             paste(r[0], ' = 2%  ',
                   rho,' = ', .(legend.values['rho','2%']), '    ',
                   eta,' = ', .(legend.values['eta','2%'])))),
           as.expression(bquote(
             paste(r[0], ' = 1.5%  ',
                   rho,' = ', .(legend.values['rho','1.5%']), '    ',
                   eta,' = ', .(legend.values['eta','1.5%']))))
         )
  )
  
  # Newell Pizer Levels
  plot.ces(output['LVL','5%'][[1]], maxbuff=0.04, col=3, main='Newell Pizer (2003) Levels Fits')
  plot.ces(output['LVL','3%'][[1]], addline=TRUE, col=1)
  plot.ces(output['LVL','2%'][[1]], addline=TRUE, col=4)
  plot.ces(output['LVL','1.5%'][[1]], addline=TRUE, col=2)
  
  abline(h=0)
  legend.values = sapply(output['LVL', c('5%','3%','2%','1.5%')], function(x) c(rho=x$rho, eta=x$eta))
  constr.rho = which(legend.values['rho',]==0)
  legend.values['rho',] = round(legend.values['rho',],3)*100
  legend.values['eta',] = round(legend.values['eta',],2)
  legend.values['rho',] = paste0(legend.values['rho',],'%')
  legend.values[,constr.rho] = paste0(legend.values[,constr.rho],'*')

  legend('top', fill=c(3,1,4,2), border=NA, bg='white', cex=1.5,
         legend = c(
           as.expression(bquote(
             paste(r[0], ' = 5%  ',
                   rho,' = ', .(legend.values['rho','5%']), '    ',
                   eta,' = ', .(legend.values['eta','5%'])))),
           as.expression(bquote(
             paste(r[0], ' = 3%  ',
                   rho,' = ', .(legend.values['rho','3%']), '    ',
                   eta,' = ', .(legend.values['eta','3%'])))),
           as.expression(bquote(
             paste(r[0], ' = 2%  ',
                   rho,' = ', .(legend.values['rho','2%']), '    ',
                   eta,' = ', .(legend.values['eta','2%'])))),
           as.expression(bquote(
             paste(r[0], ' = 1.5%  ',
                   rho,' = ', .(legend.values['rho','1.5%']), '    ',
                   eta,' = ', .(legend.values['eta','1.5%']))))
         )
  )
  
  # Newell Pizer Random Walk
  plot.ces(output['RW','5%'][[1]], maxbuff=0.04, col=3, main='Newell Pizer (2003) RW Fits')
  plot.ces(output['RW','3%'][[1]], addline=TRUE, col=1)
  plot.ces(output['RW','2%'][[1]], addline=TRUE, col=4)
  plot.ces(output['RW','1.5%'][[1]], addline=TRUE, col=2)
  
  abline(h=0)
  legend.values = sapply(output['RW', c('5%','3%','2%','1.5%')], function(x) c(rho=x$rho, eta=x$eta))
  constr.rho = which(legend.values['rho',]==0)
  legend.values['rho',] = round(legend.values['rho',],3)*100
  legend.values['eta',] = round(legend.values['eta',],2)
  legend.values['rho',] = paste0(legend.values['rho',],'%')
  legend.values[,constr.rho] = paste0(legend.values[,constr.rho],'*')
  
  legend('top', fill=c(3,1,4,2), border=NA, bg='white', cex=1.5,
         legend = c(
           as.expression(bquote(
             paste(r[0], ' = 5%  ',
                   rho,' = ', .(legend.values['rho','5%']), '    ',
                   eta,' = ', .(legend.values['eta','5%'])))),
           as.expression(bquote(
             paste(r[0], ' = 3%  ',
                   rho,' = ', .(legend.values['rho','3%']), '    ',
                   eta,' = ', .(legend.values['eta','3%'])))),
           as.expression(bquote(
             paste(r[0], ' = 2%  ',
                   rho,' = ', .(legend.values['rho','2%']), '    ',
                   eta,' = ', .(legend.values['eta','2%'])))),
           as.expression(bquote(
             paste(r[0], ' = 1.5%  ',
                   rho,' = ', .(legend.values['rho','1.5%']), '    ',
                   eta,' = ', .(legend.values['eta','1.5%']))))
         )
  )
  
  # Newell Pizer Mean Reverting
  plot.ces(output['MR','5%'][[1]], maxbuff=0.04, col=3, main='Newell Pizer (2003) MR Fits')
  plot.ces(output['MR','3%'][[1]], addline=TRUE, col=1)
  plot.ces(output['MR','2%'][[1]], addline=TRUE, col=4)
  plot.ces(output['MR','1.5%'][[1]], addline=TRUE, col=2)
  
  abline(h=0)
  legend.values = sapply(output['MR', c('5%','3%','2%','1.5%')], function(x) c(rho=x$rho, eta=x$eta))
  constr.rho = which(legend.values['rho',]==0)
  legend.values['rho',] = round(legend.values['rho',],3)*100
  legend.values['eta',] = round(legend.values['eta',],2)
  legend.values['rho',] = paste0(legend.values['rho',],'%')
  legend.values[,constr.rho] = paste0(legend.values[,constr.rho],'*')
  
  legend('top', fill=c(3,1,4,2), border=NA, bg='white', cex=1.5,
         legend = c(
           as.expression(bquote(
             paste(r[0], ' = 5%  ',
                   rho,' = ', .( format( legend.values['rho','5%'], nsmall=1 ) ), '    ',
                   eta,' = ', .( format( legend.values['eta','5%'], nsmall=2 ) )
             )
           )),
           as.expression(bquote(
             paste(r[0], ' = 3%  ',
                   rho,' = ', .(legend.values['rho','3%']), '    ',
                   eta,' = ', .(legend.values['eta','3%'])))),
           as.expression(bquote(
             paste(r[0], ' = 2%  ',
                   rho,' = ', .(legend.values['rho','2%']), '    ',
                   eta,' = ', .(legend.values['eta','2%'])))),
           as.expression(bquote(
             paste(r[0], ' = 1.5%  ',
                   rho,' = ', .(legend.values['rho','1.5%']), '    ',
                   eta,' = ', .(legend.values['eta','1.5%']))))
         )
  )
  
  # Groom et al. 2007 Regime Switching
  plot.ces(output['SS','5%'][[1]], maxbuff=0.04, col=3, main='Groom et al. (2007) SS Fits')
  plot.ces(output['SS','3%'][[1]], addline=TRUE, col=1)
  plot.ces(output['SS','2%'][[1]], addline=TRUE, col=4)
  plot.ces(output['SS','1.5%'][[1]], addline=TRUE, col=2)
  
  abline(h=0)
  legend.values = sapply(output['SS', c('5%','3%','2%','1.5%')], function(x) c(rho=x$rho, eta=x$eta))
  constr.rho = which(legend.values['rho',]==0)
  legend.values['rho',] = round(legend.values['rho',],3)*100
  legend.values['eta',] = round(legend.values['eta',],2)
  legend.values['rho',] = paste0(legend.values['rho',],'%')
  legend.values[,constr.rho] = paste0(legend.values[,constr.rho],'*')
  
  legend('top', fill=c(3,1,4,2), border=NA, bg='white', cex=1.5,
         legend = c(
           as.expression(bquote(
             paste(r[0], ' = 5%  ',
                   rho,' = ', .( format( legend.values['rho','5%'], nsmall=1 ) ), '    ',
                   eta,' = ', .( format( legend.values['eta','5%'], nsmall=2 ) )
             )
           )),
           as.expression(bquote(
             paste(r[0], ' = 3%  ',
                   rho,' = ', .(legend.values['rho','3%']), '    ',
                   eta,' = ', .(legend.values['eta','3%'])))),
           as.expression(bquote(
             paste(r[0], ' = 2%  ',
                   rho,' = ', .(legend.values['rho','2%']), '    ',
                   eta,' = ', .(legend.values['eta','2%'])))),
           as.expression(bquote(
             paste(r[0], ' = 1.5%  ',
                   rho,' = ', .(legend.values['rho','1.5%']), '    ',
                   eta,' = ', .(legend.values['eta','1.5%']))))
         )
  )
  
  # Freeman et al. 2015 Fisher Effect
  plot.ces(output['AADL','5%'][[1]], maxbuff=0.04, col=3, main='Freeman et al. (2015) AADL Fits')
  plot.ces(output['AADL','3%'][[1]], addline=TRUE, col=1)
  plot.ces(output['AADL','2%'][[1]], addline=TRUE, col=4)
  plot.ces(output['AADL','1.5%'][[1]], addline=TRUE, col=2)
  
  abline(h=0)
  legend.values = sapply(output['AADL', c('5%','3%','2%','1.5%')], function(x) c(rho=x$rho, eta=x$eta))
  constr.rho = which(legend.values['rho',]==0)
  legend.values['rho',] = round(legend.values['rho',],3)*100
  legend.values['eta',] = round(legend.values['eta',],2)
  legend.values['rho',] = paste0(legend.values['rho',],'%')
  legend.values[,constr.rho] = paste0(legend.values[,constr.rho],'*')
  
  legend('top', fill=c(3,1,4,2), border=NA, bg='white', cex=1.5,
         legend = c(
           as.expression(bquote(
             paste(r[0], ' = 5%  ',
                   rho,' = ', .( format( legend.values['rho','5%'], nsmall=1 ) ), '    ',
                   eta,' = ', .( format( legend.values['eta','5%'], nsmall=2 ) )
             )
           )),
           as.expression(bquote(
             paste(r[0], ' = 3%  ',
                   rho,' = ', .(legend.values['rho','3%']), '    ',
                   eta,' = ', .(legend.values['eta','3%'])))),
           as.expression(bquote(
             paste(r[0], ' = 2%  ',
                   rho,' = ', .(legend.values['rho','2%']), '    ',
                   eta,' = ', .(legend.values['eta','2%'])))),
           as.expression(bquote(
             paste(r[0], ' = 1.5%  ',
                   rho,' = ', .(legend.values['rho','1.5%']), '    ',
                   eta,' = ', .(legend.values['eta','1.5%']))))
         )
  )
  dev.off()
  
}


# Indeed, using an $\eta$ value of 2 (while choosing $\rho$ to achieve a 3\% near term rate) 
# would result in negative certainty-equivalent discount rates for all years after 2260.
calc_scc_ce(eta=2, target.ce=ce_rates[,'BR','3%'])$`r.ce`[which(r.ce<0), .(year, r.ce)]

# For the 3\% near-term rate, the BR term structure has the best quality of fit of all the term structure options from the literature, with an R$^2$ value of 0.995; versus 0.79-0.98 for the range of alternatives.
sapply(output[,'3%'], function(x) x$r.ce[, cor(r.ce, r.ce.target)])^2 
